#include "checksum.h"
#include <netinet/tcp.h>
#include <netinet/ip.h>
#include <string.h>

// 计算校验和的简单实现，第一个元素为ip首部
unsigned short ip_checksum(unsigned short *buf, int len) {
    unsigned long sum = 0;
    while (len > 1) {
        sum += *buf++;
        len -= 2;
    }
    if (len) {
        sum += *(unsigned char *)buf;
    }
    while (sum >> 16) {
        sum = (sum & 0xffff) + (sum >> 16);
    }
    return ~sum;
}

// 计算校验和的辅助函数
unsigned short csum(unsigned short *ptr, int nbytes) {
    long sum = 0;
    unsigned short oddbyte;
    short answer;

    while (nbytes > 1) {
        sum += *ptr++;
        nbytes -= 2;
    }
    if (nbytes == 1) {
        oddbyte = 0;
        *((u_char *)&oddbyte) = *(u_char *)ptr;
        sum += oddbyte;
    }
    sum = (sum >> 16) + (sum & 0xffff);
    sum = sum + (sum >> 16);
    answer = (short)~sum;
    return answer;
}

void tcp_checksum(iphdr *ipHeader, tcphdr *tcpHeader, char *data, int dataLen)
{
    int tcpLen = sizeof(struct tcphdr) + dataLen;
    int totalLen = tcpLen + sizeof(struct iphdr);

    // 创建伪头部
    char pseudoHeader[12];
    memset(pseudoHeader, 0, sizeof(pseudoHeader));

    // 填充伪头部
    memcpy(pseudoHeader, &ipHeader->saddr, 4);
    memcpy(pseudoHeader + 4, &ipHeader->daddr, 4);
    pseudoHeader[8] = 0;
    pseudoHeader[9] = ipHeader->protocol;
    *(short *)(pseudoHeader + 10) = htons(tcpLen);

    // 创建包含伪头部、TCP头部和数据的缓冲区
    char tcpSegment[BUFFSIZE];
    memset(tcpSegment, 0, sizeof(tcpSegment));
    memcpy(tcpSegment, pseudoHeader, 12);
    memcpy(tcpSegment + 12, tcpHeader, sizeof(struct tcphdr));
    memcpy(tcpSegment + 12 + sizeof(struct tcphdr), data, dataLen);

    // 计算校验和
    tcpHeader->check = csum((unsigned short *)tcpSegment, 12 + sizeof(struct tcphdr) + dataLen);
}
